# %%
from scipy.sparse import load_npz
import torch
from torch import optim
import torch.nn as nn
import numpy as np
import math
import torch.optim as optim
from time import time


def RMSE(Xpred, Xtest):
    mask = Xtest != 0
    Xpred = Xpred * mask
    # print(torch.sum(mask))
    error = (Xpred - Xtest) ** 2
    # print(error[:5, :5])
    return torch.sum((Xpred - Xtest) ** 2) / torch.sum(mask)


def matrix_factor(Xtrain, Xtest, k, l, max_iter, f=None, thd=100):
    """matrix factor X = UV^T using pytorch.
    Args:
        Xtrain (np.ndarray): size [user, movie].
        Xtest (np.ndarray): [user, movie].
        k (int): U [user, k], V [movie, k].
        l (float): regular param.
        mat_iter (int): max loop iterations.
        f (file): redirect output to file.
    Returns:
        U (np.ndarray).
        V (np.ndarray).
        record (list [(loss, rmse on test data)]).
    """
    train_mask = Xtrain != 0
    test_mask = Xtest != 0
    record = []

    Xtrain = torch.Tensor(Xtrain)
    Xtest = torch.Tensor(Xtest)
    mask = torch.Tensor(train_mask)
    test_mask = torch.Tensor(test_mask)
    print(train_mask.sum())

    U = torch.rand((Xtrain.shape[0], k))
    V = torch.rand((Xtrain.shape[1], k))
    U = (U - 0.5) / math.sqrt(k)
    V = (V - 0.5) / math.sqrt(k)
    U = U.requires_grad_()
    V = V.requires_grad_()

    criterion = nn.MSELoss(reduction="sum")
    optimizer = optim.Adam([U, V], lr=0.1, weight_decay=2*l)

    for i in range(max_iter):
        optimizer.zero_grad()
        Xpred = torch.mm(U, V.T)
        # print(Xpred.size())

        loss = criterion(Xpred * mask, Xtrain)
        # rmse = criterion(Xpred * test_mask, Xtest).item() / test_num
        # rmse = math.sqrt(rmse)
        rmse = RMSE(Xtest, Xpred * test_mask)
        J = 0.5 * loss.item() + l * (torch.sum(U ** 2) + torch.sum(V ** 2))
        print(f"iter {i}: loss {J}, rmse {rmse}", file=f)
        record.append((J, rmse))

        loss.backward()
        optimizer.step()
       
        record.append((J, rmse))
        if len(record) >= 20:
            cond = [record[i-1][0] - record[i][0] for i in range(5)]
            cond = np.array(cond)
            cond0 = np.all(cond > 0)
            cond1 = np.all(cond < thd)
            if cond0 and cond1:
                break


        
    return U, V, record


def RMSE(Xtest, Xpred):
    s = torch.sum((Xpred - Xtest) ** 2)
    # print(torch.sum(Xpred != 0), torch.sum(Xtest != 0))
    rmse = math.sqrt(s / torch.sum(Xtest != 0))
    return rmse
# # %%
# Xtrain = np.array(load_npz("train.npz").todense())
# Xtest = np.array(load_npz("test.npz").todense())
# k = 50
# l2 = 0.01

# # %%
# s = time()
# with open(f"torch_{k}_{l2}.txt", "w+") as f:
#     U, V, rec = matrix_factor(Xtrain, Xtest, k, l2, 100)
# print(f"cost time {time() - s} s")
# %%
